import os
import sqlite3
import hashlib
from pathlib import Path
import argparse
import zipfile
import tempfile

def check_and_connect_db(db_path):
    """检查并连接到SQLite数据库"""
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    return conn, cursor

def fetch_existing_hashes(conn, cursor):
    """从photo_hash表中获取所有记录的哈希值"""
    cursor.execute('SELECT file_hash FROM photo_hash')
    existing_hashes = {row[0] for row in cursor.fetchall()}
    return existing_hashes

def calculate_file_hash(file_path):
    """计算文件的SHA256哈希值"""
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()

def extract_zip_recursively(zip_file_path, temp_dir):
    """递归地提取ZIP文件及其内部的ZIP文件到临时目录"""
    extracted_files = []
    
    def _extract_zip(_zip_file_path, _temp_dir):
        try:
            with zipfile.ZipFile(_zip_file_path, 'r') as zip_ref:
                zip_ref.extractall(_temp_dir)
                for zip_info in zip_ref.infolist():
                    if not zip_info.is_dir():
                        inner_file_path = Path(_temp_dir) / zip_info.filename
                        if inner_file_path.suffix.lower() == '.zip':
                            # 递归处理嵌套的ZIP文件
                            nested_temp_dir = Path(tempfile.mkdtemp(dir=_temp_dir))
                            _extract_zip(inner_file_path, nested_temp_dir)
                            extracted_files.extend(extract_zip_recursively(inner_file_path, nested_temp_dir))
                        else:
                            extracted_files.append((zip_file_path, zip_info.filename, inner_file_path))
        except zipfile.BadZipFile:
            print(f"警告: 无法打开ZIP文件 {_zip_file_path}")
    
    _extract_zip(zip_file_path, temp_dir)
    return extracted_files

def scan_files(scan_dir):
    """扫描指定目录及其子目录下的所有文件和ZIP文件"""
    scan_dir_path = Path(scan_dir)
    files = []
    temp_dir = Path(tempfile.mkdtemp())
    
    for item in scan_dir_path.rglob('*'):
        if item.is_file():
            if item.suffix.lower() == '.zip':
                extracted_files = extract_zip_recursively(item, temp_dir)
                files.extend(extracted_files)
            else:
                files.append((None, item.name, item))
    
    return files

def main(db_directory, scan_dir):
    db_path = str(Path(db_directory) / 'file_hashes.db')
    conn, cursor = check_and_connect_db(db_path)

    existing_hashes = fetch_existing_hashes(conn, cursor)
    files = scan_files(scan_dir)

    for zip_file_path, inner_file_name, actual_file_path in files:
        file_hash = calculate_file_hash(actual_file_path)
        if file_hash not in existing_hashes:
            if zip_file_path:
                print(f"未找到哈希值对应的记录: ZIP文件 {zip_file_path}, 内部文件名 {inner_file_name}")
            else:
                print(f"未找到哈希值对应的记录: 文件路径 {actual_file_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="扫描指定目录中的所有文件，计算其哈希值，并检查是否存在于photo_hash表中")
    parser.add_argument('--db', required=True, help="SQLite数据库存放位置")
    parser.add_argument('--scan-dir', required=True, help="扫描目录")

    args = parser.parse_args()

    db_directory = args.db
    scan_dir = args.scan_dir

    main(db_directory, scan_dir)